#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

class KernelFastGeluGrad {
 public:
  __aicore__ inline KernelFastGeluGrad() {}
  __aicore__ inline void Init(GM_ADDR dy, GM_ADDR x, GM_ADDR z, 
                              uint32_t blockLength,
                              uint32_t tileNum, uint32_t tileLength,
                              uint32_t lasttileLength, uint32_t tilingKey) {
    ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

    this->blockLength = blockLength;
    this->tileNum =
        tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
    this->tileLength = tileLength / BUFFER_NUM;
    this->lasttileLength = lasttileLength;

    dyGm.SetGlobalBuffer((__gm__ DTYPE_X*)dy + this->blockLength * GetBlockIdx(),
                        this->blockLength);
    xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + this->blockLength * GetBlockIdx(),
                        this->blockLength);
    zGm.SetGlobalBuffer((__gm__ DTYPE_X*)z + this->blockLength * GetBlockIdx(),
                        this->blockLength);
  

    pipe.InitBuffer(inQueueIN, BUFFER_NUM, this->tileLength * 2 * sizeof(DTYPE_X));
    pipe.InitBuffer(outQueueOUT, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
    pipe.InitBuffer(tmpBuffer1, this->tileLength * sizeof(float));
    pipe.InitBuffer(tmpBuffer2, this->tileLength * sizeof(float));
    pipe.InitBuffer(tmpBuffer3, this->tileLength * sizeof(float));
    pipe.InitBuffer(tmpBuffer4, this->tileLength * sizeof(float));
  }
  __aicore__ inline void Process() {
    int32_t loopCount = this->tileNum * BUFFER_NUM;
    for (int32_t i = 0; i < loopCount; i++) {
      CopyIn(i);
      Compute(i, (DTYPE_X)0);
      CopyOut(i);
    }
  }

 private:
  __aicore__ inline void CopyIn(int32_t progress) {
    LocalTensor<DTYPE_X> inLocal = inQueueIN.AllocTensor<DTYPE_X>();

    if (BUFFER_NUM == 1) {
      if (progress == this->tileNum - 1) {
        if (progress == 0) {
          DataCopy(inLocal[0], dyGm[0], this->tileLength);
          DataCopy(inLocal[this->tileLength], xGm[0], this->tileLength);
        } else {
          DataCopy(
              inLocal[0],
              dyGm[(progress - 1) * this->tileLength + this->lasttileLength],
              this->tileLength);
          DataCopy(
              inLocal[this->tileLength],
              xGm[(progress - 1) * this->tileLength + this->lasttileLength],
              this->tileLength);
        }
      } else {
        DataCopy(inLocal[0], dyGm[progress * this->tileLength],
                 this->tileLength);
        DataCopy(inLocal[this->tileLength], xGm[progress * this->tileLength],
                 this->tileLength);
      }
    }
    if (BUFFER_NUM == 2) {
      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        DataCopy(
            inLocal[0],
            dyGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            (this->tileLength));
        DataCopy(
            inLocal[this->tileLength],
            xGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            (this->tileLength));

      }

      else {
        DataCopy(inLocal[0], dyGm[progress * (this->tileLength)],
                 (this->tileLength));
        DataCopy(inLocal[this->tileLength], xGm[progress * this->tileLength],
                 this->tileLength);
      }
    }

    inQueueIN.EnQue(inLocal);
  }
  __aicore__ inline void Compute(int32_t progress, float flag) {
    LocalTensor<float> inLocal = inQueueIN.DeQue<float>();
    LocalTensor<float> x1Local = inLocal;
    LocalTensor<float> x2Local = inLocal[this->tileLength];
    LocalTensor<float> temp1 = tmpBuffer1.Get<float>();
    LocalTensor<float> temp2 = tmpBuffer2.Get<float>();
    LocalTensor<float> temp3 = tmpBuffer3.Get<float>();
    LocalTensor<float> temp4 = tmpBuffer4.Get<float>();

    LocalTensor<float> yLocal = outQueueOUT.AllocTensor<float>();

    float attr = 1.702;
    float attr_opp = -1.702;
    //Way 1: 
    // Abs(temp1, x2Local, this->tileLength);
    // Muls(temp2, temp1, attr_opp, this->tileLength);
    // Exp(temp2, temp2, this->tileLength);

    // Mul(temp3, x2Local, temp2, this->tileLength);
    // Muls(temp3, temp3, attr, this->tileLength);

    // Sub(temp4, x2Local, temp1, this->tileLength);
    // Muls(temp4, temp4, attr, this->tileLength);
    // Exp(temp4, temp4, this->tileLength);
    
    // Add(temp1, temp2, temp3, this->tileLength);
    // Add(temp1, temp1, temp4, this->tileLength);

    // Adds(temp3, temp2, (DTYPE_X)1.0, this->tileLength);
    // Mul(temp2, temp3, temp3, this->tileLength);

    // Div(yLocal, temp1, temp2, this->tileLength);
    // Mul(yLocal, x1Local, yLocal, this->tileLength);

    //Way 2:
    Muls(temp1, x2Local, attr, this->tileLength);
    Muls(temp2, x2Local, attr_opp, this->tileLength);
    Exp(temp3, temp2, this->tileLength);
    Adds(temp3, temp3, (float)1.0, this->tileLength);
    
    Adds(temp4, temp1, (float)1.0, this->tileLength);
    Div(temp4, temp4, temp3, this->tileLength);

    Mul(temp3, temp3, temp3, this->tileLength);
    Div(temp1, temp1, temp3, this->tileLength);

    Sub(yLocal, temp4, temp1, this->tileLength);
    Mul(yLocal, yLocal, x1Local, this->tileLength);

    outQueueOUT.EnQue<float>(yLocal);

    inQueueIN.FreeTensor(inLocal);
  }

  __aicore__ inline void Compute(int32_t progress, half flag) {
    LocalTensor<half> inLocal = inQueueIN.DeQue<half>();
    LocalTensor<half> x1Local = inLocal;
    LocalTensor<half> x2Local = inLocal[this->tileLength];
    LocalTensor<float> temp1 = tmpBuffer1.Get<float>();
    LocalTensor<float> temp2 = tmpBuffer2.Get<float>();
    LocalTensor<float> temp3 = tmpBuffer3.Get<float>();
    LocalTensor<float> temp4 = tmpBuffer4.Get<float>();

    LocalTensor<half> yLocal = outQueueOUT.AllocTensor<half>();

    float attr = 1.702;
    float attr_opp = -1.702;
    //Way 2:
    Cast(temp4, x2Local, RoundMode::CAST_NONE, this->tileLength);
    Muls(temp1, temp4, attr, this->tileLength);
    Muls(temp2, temp4, attr_opp, this->tileLength);
    Exp(temp3, temp2, this->tileLength);
    Adds(temp3, temp3, (float)1.0, this->tileLength);
    
    Adds(temp4, temp1, (float)1.0, this->tileLength);
    Div(temp4, temp4, temp3, this->tileLength);

    Mul(temp3, temp3, temp3, this->tileLength);
    Div(temp1, temp1, temp3, this->tileLength);

    Cast(temp3, x1Local, RoundMode::CAST_NONE, this->tileLength);
    Sub(temp2, temp4, temp1, this->tileLength);
    Mul(temp2, temp2, temp3, this->tileLength);

    Cast(yLocal, temp2, RoundMode::CAST_NONE, this->tileLength);
    outQueueOUT.EnQue<half>(yLocal);

    inQueueIN.FreeTensor(inLocal);
  }
  __aicore__ inline void CopyOut(int32_t progress) {
    LocalTensor<DTYPE_X> outLocal = outQueueOUT.DeQue<DTYPE_X>();

    if (BUFFER_NUM == 1) {
      if (progress == this->tileNum - 1) {
        if (progress == 0) {
          DataCopy(zGm[0], outLocal, this->tileLength);
        } else {
          DataCopy(
              zGm[(progress - 1) * this->tileLength + this->lasttileLength],
              outLocal, this->tileLength);
        }
      } else {
        DataCopy(zGm[progress * this->tileLength], outLocal,
                 this->tileLength);
      }
    }
    if (BUFFER_NUM == 2) {
      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        DataCopy(
            zGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            outLocal, (this->tileLength));
      }

      else {
        DataCopy(zGm[progress * (this->tileLength)], outLocal,
                 (this->tileLength));
      }
    }

    outQueueOUT.FreeTensor(outLocal);
  }

 private:
  TPipe pipe;
  TQue<QuePosition::VECIN, BUFFER_NUM> inQueueIN;
  TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOUT;
  GlobalTensor<DTYPE_X> dyGm;
  GlobalTensor<DTYPE_X> xGm;
  GlobalTensor<DTYPE_X> zGm;

  TBuf<> tmpBuffer1;
  TBuf<> tmpBuffer2;
  TBuf<> tmpBuffer3;
  TBuf<> tmpBuffer4;
  TBuf<> tmpBuffer5;

  uint32_t blockLength;
  uint32_t tileNum;
  uint32_t tileLength;
  uint32_t lasttileLength;
};



extern "C" __global__ __aicore__ void fast_gelu_grad(GM_ADDR dy, GM_ADDR x, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelFastGeluGrad op;

    uint32_t tilingKey = 1;
    if (TILING_KEY_IS(1)) {
        tilingKey = 1;
    } else if (TILING_KEY_IS(2)) {
        tilingKey = 2;
    } else {
        tilingKey = 1;
    }

    op.Init(dy, x, z, tiling_data.blockLength,
            tiling_data.tileNum, tiling_data.tileLength,
            tiling_data.lasttileLength, 
            tilingKey);
    op.Process();
}